from typing import ClassVar, override

import matplotlib.pyplot as plt
from pyecharts import options
from pyecharts.charts import WordCloud as HTMLWordCloud
from wordcloud import WordCloud as ImageWordCloud

from ..processors import KeywordStage
from ..utils import OutputManager
from .interface import Marshaler

__all__ = ["ImageWordCloudMarshaler", "HTMLWordCloudMarshaler"]


class ImageWordCloudMarshaler(Marshaler[KeywordStage]):
    """
    Word cloud marshaler that marshals a `KeywordStage` into an image.
    """

    OUTPUT_NAME: ClassVar[str] = "wordcloud.png"

    _output_manager: OutputManager
    _config: dict

    def __init__(self, output_manager: OutputManager, global_config: dict) -> None:
        self._output_manager = output_manager
        self._config = global_config["wordcloud"]["image"]

    @override
    def marshal(self, stage: KeywordStage) -> None:
        frequencies = {keyword.word: keyword.importance for keyword in stage.keywords}
        wordcloud = ImageWordCloud(
            font_path=self._config["font_path"],
            width=self._config["width"],
            height=self._config["height"],
            background_color=self._config["background_color"],
        ).generate_from_frequencies(frequencies)

        plt.figure(figsize=(self._config["figure_width"], self._config["figure_height"]))
        plt.imshow(wordcloud, interpolation="bilinear")
        plt.axis("off")
        plt.tight_layout()
        plt.savefig(self._output_manager.filepath(ImageWordCloudMarshaler.OUTPUT_NAME), dpi=300, bbox_inches="tight")
        plt.close()


class HTMLWordCloudMarshaler(Marshaler[KeywordStage]):
    """
    Word cloud marshaler that marshals a `KeywordStage` into an HTML.
    """

    OUTPUT_NAME: ClassVar[str] = "wordcloud.html"

    _output_manager: OutputManager
    _config: dict

    def __init__(self, output_manager: OutputManager, global_config: dict) -> None:
        self._output_manager = output_manager
        self._config = global_config["wordcloud"]["html"]

    @override
    def marshal(self, stage: KeywordStage) -> None:
        wordcloud = HTMLWordCloud(
            options.InitOpts(
                width=self._config["width"],
                height=self._config["height"],
                theme=self._config["theme"],
            )
        )
        wordcloud.add(
            series_name="Keywords",
            data_pair=[(keyword.word, keyword.importance) for keyword in stage.keywords],
            shape=self._config["shape"],
        )
        wordcloud.render(str(self._output_manager.filepath(HTMLWordCloudMarshaler.OUTPUT_NAME)))
